import fitz  # PyMuPDF
import os
from PIL import Image, ImageEnhance
import numpy as np
from scipy import ndimage

def denoise_image(img, radius=1):
    """
    使用中值滤波对图像进行去噪
    :param img: PIL Image对象
    :param radius: 滤波半径，越大去噪效果越强，但可能会模糊图像
    :return: 去噪后的PIL Image对象
    """
    # 转换为numpy数组
    img_array = np.array(img)
    
    # 对每个颜色通道分别进行中值滤波
    r = ndimage.median_filter(img_array[:, :, 0], size=2*radius+1)
    g = ndimage.median_filter(img_array[:, :, 1], size=2*radius+1)
    b = ndimage.median_filter(img_array[:, :, 2], size=2*radius+1)
    
    # 合并通道
    denoised_array = np.stack([r, g, b], axis=2)
    
    # 转换回PIL Image
    return Image.fromarray(denoised_array.astype(np.uint8))

def enhance_image(image_path, color_factor=1.5, contrast_factor=1.2, denoise=False, denoise_radius=1):
    """
    增强图片的颜色饱和度和对比度，处理接近纯白和纯黑的像素，可选去噪
    :param image_path: 图片路径
    :param color_factor: 颜色饱和度增强因子（>1 增加饱和度，<1 降低饱和度）
    :param contrast_factor: 对比度增强因子
    :param denoise: 是否进行去噪处理
    :param denoise_radius: 去噪的滤波半径
    """
    # 打开图片
    img = Image.open(image_path)
    
    # 转换为RGB模式（如果不是的话）
    if img.mode != 'RGB':
        img = img.convert('RGB')
    
    
    # 获取像素数据
    pixels = img.load()
    width, height = img.size
    
    # 处理接近纯白和纯黑的像素
    threshold = 20  # 10%的255约为25
    for x in range(width):
        for y in range(height):
            r, g, b = pixels[x, y]
            
            # 检查是否接近纯白（所有通道值都接近255）
            if all(abs(255 - c) <= threshold for c in (r, g, b)):
                pixels[x, y] = (255, 255, 255)  # 设为纯白
            
            # 检查是否接近纯黑（所有通道值都接近0）
            elif all(c <= threshold for c in (r, g, b)):
                pixels[x, y] = (0, 0, 0)  # 设为纯黑
    
    # 如果启用去噪，在其他处理之前先进行去噪
    if denoise:
        img = denoise_image(img, denoise_radius)

    # 增强颜色饱和度
    color_enhancer = ImageEnhance.Color(img)
    img = color_enhancer.enhance(color_factor)
    
    # 增强对比度
    contrast_enhancer = ImageEnhance.Contrast(img)
    img = contrast_enhancer.enhance(contrast_factor)
    
    # 保存增强后的图片
    enhanced_path = image_path.replace('.png', '_enhanced.png')
    img.save(enhanced_path)
    return enhanced_path

def pdf_to_images(pdf_path, output_folder, dpi=300, enhance=True, denoise=False):
    """
    将PDF文件转换为图片并可选择增强颜色
    :param pdf_path: PDF文件路径
    :param output_folder: 输出图片的文件夹路径
    :param dpi: 输出图片的分辨率（默认300 DPI）
    :param enhance: 是否增强图片颜色
    :param denoise: 是否进行去噪处理
    """
    # 打开PDF文件
    pdf_document = fitz.open(pdf_path)
    total_pages = pdf_document.page_count

    # 确保输出文件夹存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    print(f"开始将PDF文件 {pdf_path} 转换为图片，总页数：{total_pages}")
    for page_num in range(total_pages):
        # 获取当前页面
        page = pdf_document.load_page(page_num)
        # 将页面转换为图片
        pix = page.get_pixmap(matrix=fitz.Matrix(dpi / 72, dpi / 72))
        # 保存图片
        image_path = os.path.join(output_folder, f"page_{page_num + 1}.png")
        pix.save(image_path)
        
        if enhance:
            enhanced_path = enhance_image(image_path, 2, 3, denoise=denoise)
            print(f"已保存增强后的图片：{enhanced_path}")
        else:
            print(f"已保存图片：{image_path}")
        

    print("转换完成！")

def images_to_pdf(image_folder, output_pdf, pattern="_enhanced.png"):
    """
    将文件夹中的图片合成为PDF
    :param image_folder: 图片所在文件夹路径
    :param output_pdf: 输出PDF文件路径
    :param pattern: 图片文件名模式，用于筛选要合成的图片
    """
    # 获取所有图片文件
    image_files = []
    for filename in os.listdir(image_folder):
        if pattern in filename:
            image_files.append(filename)
    
    if not image_files:
        print(f"在 {image_folder} 中没有找到包含 {pattern} 的图片")
        return
    
    # 使用页码数字进行排序
    def get_page_number(filename):
        # 从文件名中提取页码数字
        # 例如：从 "page_1_enhanced.png" 提取 1
        try:
            return int(filename.split('_')[1].split('.')[0])
        except:
            return 0
    
    # 按页码排序
    image_files.sort(key=get_page_number)
    
    # 打开并收集所有图片
    images = []
    for filename in image_files:
        image_path = os.path.join(image_folder, filename)
        images.append(Image.open(image_path))
    
    # 将第一张图片保存为PDF，并添加其他图片
    first_image = images[0]
    first_image.save(output_pdf, "PDF", save_all=True, append_images=images[1:])
    print(f"已将图片合成为PDF：{output_pdf}")
    print(f"合并的页面顺序：{', '.join(image_files)}")

# 示例使用
if __name__ == "__main__":
    pdf_path = "example.pdf"  # 替换为你的PDF文件路径
    output_folder = "output_images"  # 输出图片的文件夹路径
    output_pdf = "enhanced_book.pdf"  # 新PDF的输出路径
    
    # 首先将PDF转换为增强的图片
    pdf_to_images(pdf_path, output_folder, enhance=True, denoise=True)
    
    # 然后将增强后的图片合成为新的PDF
    images_to_pdf(output_folder, output_pdf)